import os
import pickle
from skimage.transform import resize

import numpy as np
import random
from torch.utils.data.dataloader import default_collate
from torchvision import datasets, transforms


class MNIST_Dataset():
    def __init__(self, cfg):
        self.dataset_dir = cfg.dataset_dir              #minst数据集路径，相对路径
        self.prepro_dir = cfg.prepro_dir                #处理过后的数据文件夹路径，相对路径

        self.num_instances = cfg.num_instances          #样本量
        self.image_height = cfg.image_height            #图像长
        self.image_width = cfg.image_width              #图像宽
        self.image_channel_size = cfg.image_channel_size     #通道数

        self.transform = transforms.Compose([      #transforms.Compose()类。这个类的主要作用是串联多个图片变换的操作。
            transforms.ToTensor(),    #把PIL.Image(RGB) 或者numpy.ndarray(H x W x C) [0,255] to [-1,1]，并转化成Tensor格式。
            transforms.ToPILImage(),                          #将tensor转化为PIL图像，去掉了通道这个纬度
            transforms.Grayscale(num_output_channels=1),      #将图像转化为灰度图像
            transforms.ToTensor(),
            transforms.Normalize((0.5,), (0.5,)),             #用均值和标准差归一化tensor图像[0,1] to [-1,1]
            ])


        self.prepro_train_file = os.path.join(self.prepro_dir, str(self.num_instances), 'mnist_train.pickle')
                                                     #mnist_train.pickle文件路径
        self.prepro_test_file = os.path.join(self.prepro_dir, str(self.num_instances), 'mnist_test.pickle')
                                                     #mnist_test.pickle文件路径


        if not os.path.exists(self.dataset_dir):                   #检测minst数据集路径是否存在
            os.makedirs(self.dataset_dir)                          #不存在生成文件夹
        self.load_dataset()                                        #自定义方法，加载数据集

        if not os.path.exists(self.prepro_dir):                    #检测./prepro/MINST文件夹路径是否存在
            os.makedirs(self.prepro_dir)                           #不存在生成文件夹

        if not os.path.exists(os.path.join(self.prepro_dir, str(self.num_instances))):  #./perpro/MINST/60000
            os.makedirs(os.path.join(self.prepro_dir, str(self.num_instances)))

        if not os.path.exists(self.prepro_train_file):                 #mnist_train.pickle文件是否存在
            self.preprocess_dataset(train=True)                        #自定义方法，处理数据集
            print('Load train dataset -->', self.prepro_train_file)     #输出登陆 训练 文件路径
            with open(self.prepro_train_file, 'rb') as f:  #rb：以二进制格式打开一个文件用于只读。文件指针将会放在文件的开头，默认模式
                self.train_dataset = pickle.load(f)        #从file中读取一个字符串，并将它重构为原来的python对象：训练数据集。
        else:                                                               #否则不处理preprocess_dataset
            print('Load train dataset -->', self.prepro_train_file)
            with open(self.prepro_train_file, 'rb') as f:                   #rb
                self.train_dataset = pickle.load(f)

        if not os.path.exists(self.prepro_test_file):               #同构造训练集相似，结构相同
            self.preprocess_dataset(train=False)
            print('Load test dataset -->', self.prepro_test_file)
            with open(self.prepro_test_file, 'rb') as f:
                self.test_dataset = pickle.load(f)
        else:
            print('Load test dataset -->', self.prepro_test_file)
            with open(self.prepro_test_file, 'rb') as f:
                self.test_dataset = pickle.load(f)


    def load_dataset(self):                                                #加载MNIST数据集，
        self.raw_train_dataset = datasets.MNIST(root=self.dataset_dir,     #根目录
                                                train=True,                #是否从从training.pt创建数据集，否则来自test.pt。
                                                download=True,             #TRUE 则从互联网下载数据集，然后将其放在根目录中
                                                transform=None)           #是否接受PIL图像并返回转换后版本的函数/转换

        self.raw_test_dataset = datasets.MNIST(root=self.dataset_dir,     #加载测试集
                                               train=False,
                                               download=True,
                                               transform=None)


    def preprocess_dataset(self, train=True):  #处理MNIST数据集
        if train:
            print()
            print('Preprocess train dataset')                           #train dataset
            print()
            images = self.raw_train_dataset.data.numpy()                #获得图像数据
            labels = self.raw_train_dataset.targets.numpy()             #获得标签
        else:
            print()
            print('Preprocess test dataset')                             #test dataset
            print()
            images = self.raw_test_dataset.data.numpy()                  #获得图像数据
            labels = self.raw_test_dataset.targets.numpy()               #获得标签

        _dataset = []                                                    #_dataset数组
        for i, (img, label) in enumerate(zip(images, labels)):
            '''
            #enumerate：枚举，参数X可以是一个迭代器(iterator)或者是一个序列，start是起始计数值，默认从0开始。X可以是一个字典。
              b=[1,2,3,4,5,6]
              for i,item in enumerate(b):
              print (i, item)
               Output:
                0 1
                1 2
                2 3
                3 4
                4 5
                5 6
            #zip：zip() 函数用于将可迭代的对象作为参数，将对象中对应的元素打包成一个个元组，然后返回由这些元组组成的列表。
            >> > a = [1, 2, 3]
            >> > b = [4, 5, 6]
            >> > c = [4, 5, 6, 7, 8]
            >> > zipped = zip(a, b)  # 打包为元组的列表
            [(1, 4), (2, 5), (3, 6)]
            '''
            _dataset.append((img, label))                  #将获取得到的img+label元祖写入_dataset
        if train:
            random.shuffle(_dataset)                       #将_dataset的次序打乱

        dataset = []
        instance_idx = 0                                   #样本索引
        for i, (img, label) in enumerate(_dataset):
            img = np.expand_dims(img, axis=2)             #(2, 2) to (2, 2, 1)
            '''
            #axis=1代表扩展行，axis=0代表扩展列,axis=-1代表扩展最后一个参数
            a = np.array([[1,2],[3,5]])
            y = np.expand_dims(a, axis=2)    
            y 变成了 [           
                    [ [1], [2] ]，
                    [ [3], [5] ]
                    ]
            import numpy as np
            b = np.array([1,4])
            print(b)
            >>>[1 4]
            c = np.expand_dims(b,axis=1)
            print(c)
            >>>>[[1]
                [4]]
            d = np.expand_dims(b,axis=0)
            print(d)
            >>>> [[1 4]]
            print(b.shape,c.shape,d.shape)
            >>>> (2,) (2, 1) (1, 2)
            '''
            img = resize(img, (self.image_height, self.image_width), anti_aliasing=True)       #anti_aliasing：图形保真
            img = img.astype(np.float32)                                                       #转换为float32类型
            dataset.append((self.transform(img), label, instance_idx))                 #transform过后的图片，标签，索引存入dataset
            instance_idx += 1
            if self.num_instances <= instance_idx and train:                          #多少张图片，终止条件
                break
        self.max_num_instances = len(dataset)                                         #最大样本数
        print('The number of instances: %s' % self.max_num_instances)

        if train:
            with open(self.prepro_train_file, 'wb') as f:                            #wb：是字节写入
                pickle.dump(dataset, f)                                              #将对象dataset保存到文件file中去。
        else:
            with open(self.prepro_test_file, 'wb') as f:                             #test
                pickle.dump(dataset, f)


class BatchCollator():                                                        #不太懂在干啥
    def __init__(self, image_height, image_width, image_channel_size):
        self.image_height = image_height
        self.image_width = image_width
        self.image_channel_size = image_channel_size

    def __call__(self, batch):
        batch_padded = [b for b in batch]
        return default_collate(batch_padded)
